import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

plt.rcParams['font.sans-serif'] = ['Arial Unicode MS']
df = pd.read_csv('data/honey.csv')
df = df[df['state'] != 'United States'].copy()

yearly = df.groupby('year').agg({'numcol': 'sum', 'totalprod': 'sum'}).reset_index()
plt.figure(figsize=(8, 4))
plt.plot(yearly['year'], yearly['numcol'] / 1_000_000, marker='o', label='蜂群(百万)')
plt.plot(yearly['year'], yearly['totalprod'] / 1_000_000, marker='o', label='产量(百万磅)')
plt.xlabel('年份')
plt.legend()
plt.tight_layout()
plt.savefig('第三次作业_第二题_图1.png', dpi=300)
plt.close()

top_states = df.groupby('state')['numcol'].mean().sort_values(ascending=False).head(6).index
small = df[df['state'].isin(top_states)]
pivot = small.pivot_table(index='state', columns='year', values='yieldpercol')
plt.figure(figsize=(10, 4))
sns.heatmap(pivot, cmap='YlOrBr')
plt.ylabel('州')
plt.tight_layout()
plt.savefig('第三次作业_第二题_图2.png', dpi=300)
plt.close()

plt.figure(figsize=(6, 4))
plt.scatter(df['yieldpercol'], df['priceperlb'], alpha=0.5)
plt.xlabel('单位产量(磅/群)')
plt.ylabel('价格(美元/磅)')
plt.tight_layout()
plt.savefig('第三次作业_第二题_图3.png', dpi=300)
plt.close()
